"""Activity for generating trace summaries using LLM."""

from uuid import uuid4

import structlog
import temporalio

from posthog.schema import DateRange, TraceQuery

from posthog.hogql_queries.ai.trace_query_runner import TraceQueryRunner
from posthog.models.event.util import create_event
from posthog.models.team import Team
from posthog.sync import database_sync_to_async
from posthog.temporal.llm_analytics.trace_summarization import constants
from posthog.temporal.llm_analytics.trace_summarization.models import SummarizationActivityResult

from products.llm_analytics.backend.summarization.llm import summarize
from products.llm_analytics.backend.summarization.llm.schema import SummarizationResponse
from products.llm_analytics.backend.text_repr.formatters import (
    FormatterOptions,
    format_trace_text_repr,
    llm_trace_to_formatter_format,
)

logger = structlog.get_logger(__name__)


@temporalio.activity.defn
async def generate_and_save_summary_activity(
    trace_id: str,
    team_id: int,
    window_start: str,
    window_end: str,
    mode: str,
    batch_run_id: str,
    model: str | None = None,
) -> SummarizationActivityResult:
    """
    Generate summary for a trace and save it to ClickHouse.

    Fetches trace data, generates LLM summary, and saves it as an event.
    Saves directly to avoid passing large objects through workflow history.
    """

    def _fetch_trace_and_format(
        trace_id: str, team_id: int, window_start: str, window_end: str
    ) -> tuple[dict, list, str, Team] | None:
        """Fetch trace data and format text representation."""
        team = Team.objects.get(id=team_id)

        query = TraceQuery(
            traceId=trace_id,
            dateRange=DateRange(date_from=window_start, date_to=window_end),
        )

        runner = TraceQueryRunner(team=team, query=query)
        response = runner.calculate()

        if not response.results:
            return None  # Trace not found in window

        llm_trace = response.results[0]
        trace_dict, hierarchy = llm_trace_to_formatter_format(llm_trace)

        options: FormatterOptions = {
            "include_line_numbers": True,
            "truncated": False,
            "include_markers": False,
            "collapsed": False,
        }

        text_repr = format_trace_text_repr(
            trace=trace_dict,
            hierarchy=hierarchy,
            options=options,
        )

        return trace_dict, hierarchy, text_repr, team

    def _save_summary_event(summary_result: SummarizationResponse, hierarchy: list, text_repr: str, team: Team) -> None:
        """Save summary as $ai_trace_summary event to ClickHouse."""

        event_uuid = uuid4()

        summary_bullets_json = [bullet.model_dump() for bullet in summary_result.summary_bullets]
        summary_notes_json = [note.model_dump() for note in summary_result.interesting_notes]

        properties = {
            "$ai_trace_id": trace_id,
            "$ai_batch_run_id": batch_run_id,
            "$ai_summary_mode": mode,
            "$ai_summary_title": summary_result.title,
            "$ai_summary_flow_diagram": summary_result.flow_diagram,
            "$ai_summary_bullets": summary_bullets_json,
            "$ai_summary_interesting_notes": summary_notes_json,
            "$ai_text_repr_length": len(text_repr),
            "$ai_event_count": len(hierarchy),
        }

        create_event(
            event_uuid=event_uuid,
            event=constants.EVENT_NAME_TRACE_SUMMARY,
            team=team,
            distinct_id=f"trace_summary_{team_id}",
            properties=properties,
        )

    # Fetch trace data and format text representation
    result = await database_sync_to_async(_fetch_trace_and_format, thread_sensitive=False)(
        trace_id, team_id, window_start, window_end
    )

    # Handle trace not found in window
    if result is None:
        logger.warning(
            "Skipping trace - not found in time window",
            trace_id=trace_id,
            window_start=window_start,
            window_end=window_end,
        )
        return SummarizationActivityResult(
            trace_id=trace_id,
            success=False,
            skipped=True,
            skip_reason="trace_not_found",
        )

    trace, hierarchy, text_repr, team = result

    # Check if trace is too large for LLM processing
    if len(text_repr) > constants.MAX_TEXT_REPR_LENGTH:
        logger.warning(
            "Skipping trace - text representation too large for LLM",
            trace_id=trace_id,
            text_repr_length=len(text_repr),
            max_length=constants.MAX_TEXT_REPR_LENGTH,
        )
        return SummarizationActivityResult(
            trace_id=trace_id,
            success=False,
            skipped=True,
            skip_reason="text_repr_too_large",
            text_repr_length=len(text_repr),
            event_count=len(hierarchy),
        )

    # Generate summary using LLM
    summary_result = await summarize(
        text_repr=text_repr,
        team_id=team_id,
        mode=mode,
        model=model,
    )

    # Save event to ClickHouse immediately
    await database_sync_to_async(_save_summary_event, thread_sensitive=False)(
        summary_result, hierarchy, text_repr, team
    )

    return SummarizationActivityResult(
        trace_id=trace_id,
        success=True,
        text_repr_length=len(text_repr),
        event_count=len(hierarchy),
    )
